from scipy import array, exp, log, cumsum
from Bio import Phylo, SeqIO, Entrez, AlignIO
from Bio.Align import MultipleSeqAlignment, AlignInfo
import cPickle, tempfile, re, sys, os
import weblogolib
from array import array as binarray
from copy import deepcopy

def getstrslice(s, r):
	return eval('s['+r+']')

class PhylipTree:
	
	def __init__(self, treepath, seqpath, ladderize=False, nexus_cutoff=1):
		self.tree = Phylo.read(treepath,'newick')
		if ladderize:
			self.tree.ladderize();
		self.seqs={}
		for seq_record in SeqIO.parse(seqpath, "fasta"):
			self.seqs[seq_record.id[:30]] = seq_record
		if os.access('.phylitree.pickle', os.R_OK):
			with open('.phylitree.pickle') as f:
				pickleval = cPickle.load(f)
				self.gis = pickleval['gis']
				self.taxa = pickleval['taxa']
		else:
			self.gis = {}
			self.taxa = {}
		givals = set(map(lambda x : x.id.split("|")[1], self.seqs.values())).difference(self.gis)
		for (i,gi) in enumerate(givals):
			sys.stdout.write('\r' + ' '*160)
			sys.stdout.write('\rDownload   missing protein  info for GI:' + str(gi) + ' ' + str(i+1) + ', ' + str(len(givals)-i) + ' to go')
			sys.stdout.flush()
			handle = Entrez.esummary(db="protein", id=gi)
			self.gis[gi] = Entrez.read(handle)[0]
			taxid = self.gis[gi]['TaxId']
			if(taxid not in self.taxa):
				sys.stdout.write('\rDownload   missing taxonomy info for GI:' + str(gi) + ' ' + str(i+1) + ', ' + str(len(givals)-i) + ' to go')
				sys.stdout.flush()
				handle = Entrez.efetch(db="Taxonomy", id=taxid)
				self.taxa[taxid] = Entrez.read(handle)[0]
		if len(givals):
			sys.stdout.write('\r' + ' '*160 + '\r')
			sys.stdout.flush()
			pickleval = {'gis' : self.gis, 'taxa' : self.taxa}
			with open('.phylitree.pickle', 'w') as f:
				cPickle.dump(pickleval, f)
		givals = map(lambda x : x.id.split("|")[1], self.seqs.values())
		if len(givals) < len(self.gis):
			self.gis = dict(map(lambda x : (x, self.gis[x]), givals))
		taxids = map(lambda x : x['TaxId'], self.gis.values())
		if len(taxids) < len(self.taxa):
			self.taxa = dict(map(lambda x : (x, self.taxa[x]), taxids))
		self.set_nexus_cutoff(nexus_cutoff)

	def __define_commands_(self):
		self.com = 	{
						'collapse' : 
							{'nexid' : self.collapse_nexid },
						'prune' :
							{'nexid' 	: self.prune_nexid,
							 'gi'    	: self.prune_gi ,
							 'sequence'	: self.prune_sequence,
							 'genus'	: self.prune_genus },
						'extract' :
							{'nexid'	: self.extract_nexid }
					}

	def set_nexus_cutoff(self, nexus_cutoff=1):
		self.nexus_cutoff=nexus_cutoff
		self.nexi = self.get_nexi()

	def export_sequence_pattern(self, options):
		path, gi, aalist = options.fasta_export.split(':')
		self.__export_sequence_pattern_(gi, aalist, path)

	def __export_sequence_pattern_(self, gi, aalist, path):
		gi = str(gi)
		seqptrn = self.seqs[filter(lambda x : x.split('|')[1]==gi, self.seqs.keys())[0]]
		ind = cumsum(map(lambda x : int(x!='-'), seqptrn)).tolist()
		numptrn = map(lambda x : ind.index(int(x[1:])), aalist.split(','))
		al1 = MultipleSeqAlignment(deepcopy(self.seqs.values()))
		filtered_sequences = map(lambda x : binarray('c',''.join(x)), array(al1)[:,numptrn])
		mutseqs = map(lambda x : x.seq.tomutable(), al1)
		for (i, record) in enumerate(al1):
			mutseqs[i].data = filtered_sequences[i]
			record.seq = mutseqs[i].toseq()
		AlignIO.write(al1, path, 'fasta')

	def edit(self, comfile):
		if('com' not in dir(self)):
			self.__define_commands_()
		with open(comfile) as f:
			for line in f:
				if line.strip():
					if line.strip()[0] != "#":
						self.execute_command(line.split("#")[0].split())

	def execute_command(self, arguments):
		self.com[arguments[0]][arguments[1]](arguments[2:])
		self.nexi = self.get_nexi()

	def collapse_nexid(self, arguments):
		print "Collapsing nexus #" + arguments[0]
		nexus = self.get_nexus(int(arguments[0]))
		method = self.tree.collapse
		if len(arguments)>1:
			if arguments[1] == 'all':
				method = self.tree.collapse_all
		for clade in nexus.clades:
			for node in clade.get_nonterminals():
				method(node)

	def prune_nexid(self, arguments):
		nexus = self.get_nexus(int(arguments[0]))
		method = self.tree.prune
		if len(arguments)>1:
			clade = nexus.clades[int(arguments[1])-1]
			for node in clade.get_terminals():
				print "Pruning " + str(node)
				method(node)
		else:
			for clade in nexus.clades:
				for node in clade.get_terminals():
					print "Pruning " + str(node)
					method(node)

	def extract_nexid(self, arguments):
		nexus = self.get_nexus(int(arguments[0]))
		for node in self.tree.get_terminals():
			if not nexus.is_parent_of(node):
				self.tree.prune(node)

	def extract_nexus(self, options):
		subtree = self.tree.from_clade(self.get_nexus(options.nexus))
		Phylo.write(subtree, options.extract_nexus, 'newick')

	def prune_gi(self, arguments):
		leafs = self.get_gi_leaf(arguments[0])
		for leaf in leafs:
			print "Pruning " + str(leaf)
			self.tree.prune(leaf)

	def prune_sequence(self, arguments):
		aaranges = map(lambda x : (x+[x[0]])[:2], map(lambda x : x.split(':'), arguments[1].split(',')))
		al1 = self.get_alignment(self.nexi[0])
		extation = arguments[0] == 'match'
		gis2prune = []
		for record in al1:
			ptrn = ''.join(map(lambda x : str(record.seq).__getslice__(int(x[0])-1,int(x[1])), aaranges))
			if (bool(re.match(arguments[2], ptrn))) == extation:
				gis2prune.append(record.id.split('|')[1])
		print "Pruning " + str(len(gis2prune)) + " leafs based on sequence " + arguments[0]
		for gi in gis2prune:
			for leaf in self.get_gi_leaf(gi):
				self.tree.prune(leaf)

	def prune_genus(self, arguments):
		al1 = self.get_alignment(self.nexi[0])
		extation = arguments[0] == 'match'
		gis2prune = []
		for record in al1:
			taxon = self.taxa[self.gis[record.id.split('|')[1]]['TaxId']]['ScientificName'].split()[0]
			if (bool(re.match(arguments[1], taxon))) == extation:
				gis2prune.append(record.id.split('|')[1])
		print "Pruning " + str(len(gis2prune)) + " leafs based on genus " + arguments[0]
		for gi in gis2prune:
			for leaf in self.get_gi_leaf(gi):
				self.tree.prune(leaf)

	def get_gi_leaf(self, gi):
		return filter(lambda x : x.name.split('|')[1]==gi, self.tree.get_terminals())

	def get_nexus_path(self, nexus):
		if type(nexus) == int:
			nexus = self.nexi[nexus]
		return map(self.nexi.index, filter(lambda x : x in self.nexi, self.tree.get_path(nexus)))

	def print_nexus_path(self, nexid):
		print "Show path to nexus #" + str(nexid)
		print '-'.join(map(str, self.get_nexus_path(nexid)))

	def print_nexus_paths(self):
		print "Nexus #   Split              Children   Path"
		paths = map(self.get_nexus_path, self.nexi)
		parents = array(map(lambda x : x[-2:], filter(lambda x : len(x)>1, paths)))
		for (nexid, nexus) in enumerate(self.nexi):
			path = paths[nexid]
			if parents.any():
				children = parents.T[1][parents.T[0]==nexid]
			else:
				children = []
			print "%7d | %-16s | %-8s | %s" % (nexid, clade_split_string(nexus), '/'.join(map(str,children)),'-'.join(map(str, path)))

	def nexus_branches_string(self, nexus):
		return '/'.join(map(self.get_nexid_string, nexus.clades))

	def pickle_dump(self, path):
		self.com = None
		del self.com
		with open(path, 'w') as f:
			cPickle.dump(self, f)
		self.__define_commands_()

	def nexus_draw(self, nexus):
		Phylo.draw_ascii(nexus)

	def get_nodes(self):
		return self.tree.get_nonterminals()

	def get_nexus(self, nexus):
		return self.get_nexi()[nexus]

	def get_nexid(self, nexus):
		return self.get_nexi().index(nexus)

	def get_nexid_string(self, nexus):
		if nexus in self.nexi:
			return str(self.get_nexi().index(nexus))
		else:
			return 'X'

	def get_leafs(self):
		return self.tree.get_terminals()

	def get_nexi(self):
		return filter(self.is_nexus, self.tree.get_nonterminals())

	def info_report(self):
		print "-----------------"
		print "Tree information:"
		print "%d terminals" % self.tree.count_terminals()
		print "%d non-terminals" % len(self.tree.get_nonterminals())
		print "-----------------"

	def nexus_report(self, verbose = 0):
		sys.stdout.write('Nexus #  Split')
		if verbose>0:
			sys.stdout.write("    Info content    Species number")
		print
		for (i,nexus) in enumerate(self.get_nexi()):
			line = '%7d %4d/%-4d' % (tuple([i]+get_clade_split(nexus)))
			if verbose>0:
				line += '  %5.2f/%5.2f  %4d/%-4d' % (tuple(map(self.get_info_content, nexus.clades)+map(len, map(self.get_node_geni,nexus.clades))))
			print line

	def print_nexus_aa(self, options):
		nexid, aa, degap = options.nexus, options.nexus_aa, options.degap
		nexus = self.nexi[nexid]
		if options.no_clades:
			print "Sequence info for nexus #" + str(nexid)
			self.print_node_aa(nexus, aa, degap)
		else:
			for (i, clade) in enumerate(self.nexi[nexid].clades):
				print "Sequence info for clade #" + str(nexid) + '-' + str(i+1)
				self.print_node_aa(clade, aa, degap)

	def print_node_aa(self, node, aa, degap=True):
		aaranges = map(lambda x : (x+[x[0]])[:2], map(lambda x : x.split(':'), aa.split(',')))
		al1 = self.get_alignment(node)
		for record in al1:
			gi = record.id.split('|')[1]
			taxon = self.taxa[self.gis[gi]['TaxId']]['ScientificName']
			title = self.gis[gi]['Title']
			seqslices = map(lambda x : str(record.seq).__getslice__(int(x[0])-1,int(x[1])), aaranges)
			if degap:
				print '|'.join(map(lambda x : x.replace('-','').ljust(len(x)),seqslices)), ("GI:"+gi).ljust(15), title
			else:
				print '|'.join(seqslices), ("GI:"+gi).ljust(15), title

	def get_info_content(self, node):
		return array(self.get_info_chart(node)).mean()

	def get_info_chart(self, node):
		return map(info2, self.get_ungapped_freqs(node))

	def get_ungapped_freqs(self, node):
		return map(lambda x : map(lambda t : x.count(t), set(x)), self.get_ungapped_cols(node))

	def get_ungapped_cols(self, node, aa=None):
		al1 = self.get_alignment(node)
		if aa:
			aaranges = map(lambda x : (x+[x[0]])[:2], map(lambda x : x.split(':'), aa.split(',')))
			seqslice = array(reduce(lambda x,y : x+y, map(lambda x : range(int(x[0])-1,int(x[1])), aaranges)))
			return map(''.join, array(al1)[:,seqslice].T)
		else:
			return map(''.join, array(al1)[:,(array(al1)!='-').any(0)].T)

	def get_alignment(self, node):
		return MultipleSeqAlignment(map(self.seqs.get, map(lambda x : x.name, node.get_terminals())))

	def get_node_genus_list(self, node):
		return map(self.get_leaf_genus, node.get_terminals())

	def get_node_geni(self, node):
		return sorted(set(self.get_node_genus_list(node)))

	def get_leaf_genus(self, leaf):
		return self.get_leaf_species(leaf).split()[0]

	def get_leaf_species(self, leaf):
		return self.taxa[self.gis[leaf.name.split('|')[1]]['TaxId']]['ScientificName']

	def get_ungapped_rows(self, node, aa=None):
		al1 = self.get_alignment(node)
		if aa:
			aaranges = map(lambda x : (x+[x[0]])[:2], map(lambda x : x.split(':'), aa.split(',')))
			seqslice = array(reduce(lambda x,y : x+y, map(lambda x : range(int(x[0])-1,int(x[1])), aaranges)))
			return map(''.join,array(al1)[:,seqslice])
		else:
			return map(''.join,array(al1)[:,(array(al1)!='-').any(0)])

	def get_node_species(self, node):
		return sorted(set(self.get_node_species_list(node)))

	def get_node_species_list(self, node):
		return map(self.get_leaf_species, node.get_terminals())

	def nexus_species(self, options):
		nexus = self.get_nexus(options.nexus)
		if options.no_clades:
			print '\n'.join(sorted(set(map(lambda x : ' '.join(x.split()[:2]), self.get_node_species(nexus)))))
		else:
			cladesplit = get_clade_split(nexus)
			species = map(self.get_node_species, nexus.clades)
			print "Clade 1 (" + str(cladesplit[0]) + "proteins): \n" + "   \n".join(species[0])
			print "Clade 2 (" + str(cladesplit[1]) + "proteins): \n" + "   \n".join(species[1])

	def get_leaf_lineage(self, leaf):
		gi = leaf.name.split('|')[1]
		taxon = self.taxa[self.gis[gi]['TaxId']]
		lineage = dict(map(lambda x : (x.get('Rank','no rank'), x.get('ScientificName', 'Unknown')), taxon['LineageEx']))
		lineage['species'] = lineage.get('species', taxon['ScientificName'])
		lineage['GI'] = gi
		return lineage

	def get_node_lineages(self, node):
		return map(self.get_leaf_lineage, node.get_terminals())

	def print_node_lineage(self, node, cols, fUnique=False):
		lineages = map(lambda x : map(str,map(x.get, cols)), self.get_node_lineages(node))
		if fUnique:
			lines = map('|'.join,lineages)
			lineages = sorted(set(lines))
			linecounts = map(lines.count, lineages)
			lineages = map(lambda x : x.split('|'), lineages)
		widths = map(lambda x : x+2, map(max, map(lambda x : map(len, x), zip(*([cols]+lineages)))))
		if fUnique:
			hbar=' '+'-'*(sum(widths)+len(widths)+7)+' '
			xlabel = ' Count |'
			print str(len(lineages)) + ' unique lineages found'
		else:
			hbar=' '+'-'*(sum(widths)+len(widths)-1)+' '
			xlabel = ''
		print hbar
		print '|' + '|'.join(map(lambda x: cols[x].center(widths[x]), range(len(cols)))) + '|'+xlabel
		print hbar
		if fUnique:
			countindex = sorted(range(len(linecounts)), key=linecounts.__getitem__, reverse=True)
			for i in countindex:
				print '|' + '|'.join(map(lambda x: lineages[i][x].center(widths[x]), range(len(cols)))) + '|' + str(linecounts[i]).center(7) +'|'
		else:
			for lineage in lineages:
				print '|' + '|'.join(map(lambda x: lineage[x].center(widths[x]), range(len(cols)))) + '|'
		print hbar

	def print_nexus_lineage(self, options, fUnique=False):
		nexus = self.get_nexus(options.nexus)
		cols = options.lineage_cols.split(',')
		if fUnique:
			if 'GI' in cols:
				cols.remove('GI')
		else:
			if 'GI' not in cols:
				cols.insert(0, 'GI')
		if options.no_clades:
			print "Lineage table for nexus #"+str(options.nexus)
			self.print_node_lineage(nexus, cols, fUnique)
		else:
			for (i,clade) in enumerate(nexus.clades):
				print "Lineage table for nexus #"+str(options.nexus)+" clade #"+str(i+1)
				self.print_node_lineage(clade, cols, fUnique)

	def print_nexus_unique_lineages(self, options):
		self.print_nexus_lineage(options, fUnique=True)

	def nexus_geni(self, options):
		nexus = self.get_nexus(options.nexus)
		if options.no_clades:
			print ','.join(self.get_node_geni(nexus))
		else:
			cladesplit = get_clade_split(nexus)
			geni = map(self.get_node_geni, nexus.clades)
			print "Clade 1 (" + str(cladesplit[0]) + "proteins): " + ",".join(geni[0])
			print "Clade 2 (" + str(cladesplit[1]) + "proteins): " + ",".join(geni[1])

	def print_nexi_with_genus(self, options):
		genus = options.nexi_with_genus
		for (i, nexus) in enumerate(self.nexi):
			genterms = self.get_node_genus_list(nexus)
			Ng = len(filter(lambda x : genus==x, genterms))
			if Ng:
				Nt = len(genterms)
				print "%4d %4d out of %4d" % (i, Ng, Nt)

	def nexus_write(self, options, prefix='nexus', altype='fasta'):
		nexid = options.nexus
		nexus = self.get_nexus(nexid)
		if not options.no_clades:
			for (i,node) in enumerate(nexus.clades):
				al1 = self.get_alignment(node)
				AlignIO.write(al1, '.'.join([prefix, str(nexid), str(i+1), altype]), altype)
		al1 = self.get_alignment(nexus)
		AlignIO.write(al1, '.'.join([prefix, str(nexid), altype]), altype)

	def write_weblogo(self, path, node, imgtype='pdf', aa=None, annotate=None, number_fontsize=8):
		tmpath = tempfile.mktemp()
		with open(tmpath,'w') as f:
			for line in self.get_ungapped_rows(node, aa):
				f.write(line+'\n')
		with open(tmpath) as f:
			seqs = weblogolib.read_seq_data(f)
		data = weblogolib.LogoData.from_seqs(seqs)
		options = weblogolib.LogoOptions()
		if annotate:
			options.annotate = annotate.split(',')
		options.title = "Node"
		options.number_fontsize = number_fontsize
		#options.rotate_numbers = True
		logoformat = weblogolib.LogoFormat(data, options)
		formatters = {  'png' : weblogolib.png_formatter,
						'eps' : weblogolib.eps_formatter,
						'jpg' : weblogolib.jpeg_formatter,
						'pdf' : weblogolib.pdf_formatter,
						'pngprint' : weblogolib.png_print_formatter,
						'txt' : weblogolib.txt_formatter}
		for value in imgtype.split(','):
			with open('.'.join([path,value]),'w') as f:
				formatters[value]( data, logoformat, f)

	def nexus_weblogo(self, options, prefix='nexus'):
		imgtype = options.imgtype
		nexid, aa, no_clades, annotate = options.nexus, options.nexus_aa, options.no_clades, options.annotate
		number_fontsize = options.weblogo_number_fontsize
		nexus = self.get_nexi()[nexid]
		if not no_clades:
			for (i,node) in enumerate(nexus.clades):
				path = '.'.join([prefix, str(nexid), str(i+1)])
				self.write_weblogo(	path, 
									node, 
									imgtype=imgtype, 
									aa=aa, 
									annotate=annotate, 
									number_fontsize=number_fontsize)
		self.write_weblogo('.'.join([prefix, str(nexid)]), nexus, imgtype=imgtype, aa=aa, annotate=annotate, number_fontsize=number_fontsize)

	def is_nexus(self, node):
		return min(get_clade_split(node))>self.nexus_cutoff

def get_clade_split(node):
	return map(lambda x : x.count_terminals(), node.clades)

def clade_split_string(node):
	return '/'.join(map(lambda x : str(x.count_terminals()), node.clades))

''' --------------------------- '''


def printal(node):
	al1=MultipleSeqAlignment(map(seqs.get, map(lambda x : x.name, node.get_terminals())))
	sumal1=AlignInfo.SummaryInfo(al1)
	for line in map(''.join,(array(al1)[:,(array(al1)!='-').any(0)].tolist())):
		print line

def write_alignment(path, node):
	with open(path,'w') as f:
		for line in get_ungapped_rows(node):
			f.write(line+'\n')


def get_position_letter(node, i):
	return get_alignment(node, seqs)[:,i]

def get_pssm(node):
	return AlignInfo.SummaryInfo(get_alignment(node, seqs)).pos_specific_score_matrix()

def info2(x):
	return (log(sum(x))-sum(log(x)*x)/sum(x))/log(2)

def get_clade_sizes(tree):
	return map(lambda x : x.count_terminals(), tree.get_nonterminals())

def get_branch_lengths(tree):
	return map(lambda x : x.branch_length, tree.get_nonterminals())

def get_preterminals(tree):
	return filter(lambda x : x.is_preterminal(), tree.get_nonterminals())


